import re
import os
from typing import Dict, Any, List, Optional, Union, Tuple
from collections import defaultdict


def resolve_alias_chain(name: str, aliases: Dict[str, str], visited=None) -> str:
    """递归解析别名链，防止循环引用"""
    if visited is None:
        visited = set()
    
    if name in visited:
        return name  # 检测到循环引用
    
    visited.add(name)
    
    if name in aliases:
        return resolve_alias_chain(aliases[name], aliases, visited)
    return name

def build_tree(root: str, relationships: Dict[str, List[Dict]], relationship_labels: Dict[str, List[Dict]], 
               aliases: Dict[str, str], filter_depth: Union[int, Dict[str, int]] = 2, 
               filter_mode: str = "down", exclude_patterns: List[str] = None) -> Dict[str, Any]:
    """
    构建函数调用树，显示目标函数调用了哪些函数
    
    参数:
        root: 根函数名
        relationships: 函数调用关系
        relationship_labels: 函数标签
        aliases: 函数别名映射
        filter_depth: 过滤深度，可以是整数（对所有模式使用相同深度）或字典（为每个模式指定不同深度）
                    例如：{{"down": 2, "up": 3}}，默认为2
        filter_mode: 过滤模式，'down'表示从上向下过滤，'up'表示从叶节点向上过滤，
                    'both'表示同时启用两种过滤，'none'表示不过滤，默认为'down'
        exclude_patterns: 要排除的函数名列表，这些函数将在树构建过程中被过滤掉
    """
    visited = set()
    node_has_children = {}  # 跟踪节点在首次处理期间是否在原始关系中有子节点
    
    # 处理filter_depth参数
    if isinstance(filter_depth, int):
        # 如果是整数，为所有模式使用相同的深度
        depth_settings = {"down": filter_depth, "up": filter_depth}
    else:
        # 如果是字典，使用指定的深度，对未指定的模式使用默认值2
        depth_settings = {"down": 2, "up": 2}
        depth_settings.update(filter_depth)
    
    # 处理filter_mode参数
    filter_modes = []
    if filter_mode == "both":
        filter_modes = ["down", "up"]
    elif filter_mode == "none":
        filter_modes = []  # Empty list means no filtering
    else:
        filter_modes = [filter_mode]
    
    # 预计算分支重要性
    important_nodes = set()
    original_root = resolve_alias_chain(root, aliases)
    
    # 总是认为根函数是重要的
    important_nodes.add(original_root)
    
    # If we're not filtering (mode is 'none'), mark all nodes as important
    if not filter_modes or filter_mode == "none":
        for node in relationships.keys():
            resolved_node = resolve_alias_chain(node, aliases)
            important_nodes.add(resolved_node)
        
        # Also add all called functions
        for calls in relationships.values():
            for call in calls:
                func_name = call['func']
                resolved_func = resolve_alias_chain(func_name, aliases)
                important_nodes.add(resolved_func)
        
        # Skip the rest of the filtering logic
        filter_modes = []
    
    # 从上向下过滤 - 标记重要节点
    if "down" in filter_modes:
        down_depth = depth_settings["down"]
        
        def mark_important_nodes_down(start_node: str):
            """标记重要节点及其后代"""
            visited_mark = set()
            paths = []
            
            def dfs(node, path, depth=0):
                # 解析别名到原始名称
                original_node = resolve_alias_chain(node, aliases)
                
                if original_node in visited_mark:
                    return
                
                visited_mark.add(original_node)
                current_path = path + [original_node]
                
                # 如果达到filter_depth深度，检查是否有被调用的函数
                if depth >= down_depth and original_node in relationships and relationships[original_node]:
                    # 这个节点有调用其他函数，因此它和它的所有祖先都是重要的
                    paths.append(current_path)
                    return
                
                # 继续深度优先搜索
                if original_node in relationships:
                    for call in relationships[original_node]:
                        dfs(call['func'], current_path, depth + 1)
            
            dfs(start_node, [])
            
            # 将所有重要路径中的节点标记为重要
            for path in paths:
                for node in path:
                    important_nodes.add(node)
        
        # 从根函数开始标记重要节点
        mark_important_nodes_down(original_root)
    
    # 从叶节点向上过滤 - 标记不重要的节点
    if "up" in filter_modes:
        up_depth = depth_settings["up"]
        
        # 首先构建反向调用图（被调用者 -> 调用者）
        reverse_call_graph = defaultdict(list)
        for caller, calls in relationships.items():
            for call in calls:
                callee = resolve_alias_chain(call['func'], aliases)
                reverse_call_graph[callee].append(caller)
        
        # 找出所有叶节点（没有调用其他函数的节点）
        leaf_nodes = set()
        for node in relationships:
            original_node = resolve_alias_chain(node, aliases)
            if not relationships.get(original_node, []):
                leaf_nodes.add(original_node)
        
        # 对于每个叶节点不在relationships中的节点，也将其视为叶节点
        for caller, calls in relationships.items():
            for call in calls:
                callee = resolve_alias_chain(call['func'], aliases)
                if callee not in relationships:
                    leaf_nodes.add(callee)
        
        # 从每个叶节点向上追溯up_depth层，标记这些节点为不重要
        unimportant_nodes = set()
        for leaf in leaf_nodes:
            # 标记叶节点本身
            if leaf != original_root:  # 除了根节点
                unimportant_nodes.add(leaf)
            
            # 向上追溯并标记祖先
            current_level = [leaf]
            for level in range(up_depth - 1):  # -1 是因为已经标记了叶节点
                next_level = []
                for node in current_level:
                    for parent in reverse_call_graph.get(node, []):
                        parent_original = resolve_alias_chain(parent, aliases)
                        if parent_original != original_root:  # 不标记根节点
                            unimportant_nodes.add(parent_original)
                            next_level.append(parent_original)
                
                if not next_level:
                    break
                
                current_level = next_level
        
        # 仅在单独使用'up'模式时才需要这一步
        if "down" not in filter_modes:
            # 标记所有不在unimportant_nodes中的节点为重要
            for node in relationships:
                original_node = resolve_alias_chain(node, aliases)
                if original_node not in unimportant_nodes:
                    important_nodes.add(original_node)
        else:
            # 如果同时使用'up'和'down'模式，从important_nodes中移除unimportant_nodes
            important_nodes = important_nodes - unimportant_nodes
    
    def _build(node: str) -> Optional[Dict[str, Any]]:
        # 解析别名到原始名称
        original_node = resolve_alias_chain(node, aliases)
        
        # Check if this node should be excluded based on exclude_patterns
        if exclude_patterns and original_node in exclude_patterns:
            return None
        
        # 确定节点名称（是否带特殊标记）
        current_name = f'** {original_node}' if relationship_labels and original_node in relationship_labels else original_node
        
        if original_node in visited:
            # 处理重复节点：根据原始是否有子节点添加前缀
            if node_has_children.get(original_node, False):
                return {'name': f'repeat## {current_name}', 'children': []}
            else:
                return {'name': current_name, 'children': []}
        visited.add(original_node)
        
        # 记录该节点原始是否有子节点
        original_has_children = original_node in relationships and bool(relationships.get(original_node))
        node_has_children[original_node] = original_has_children
        
        # Special case for the root function itself - always include it and all its direct calls
        if original_node == original_root:
            # Recursively build children for the root node
            children = []
            if original_node in relationships:
                for call in relationships[original_node]:
                    # Skip excluded functions
                    if exclude_patterns and resolve_alias_chain(call['func'], aliases) in exclude_patterns:
                        continue
                    child_node = _build(call['func'])
                    if child_node is not None:
                        child_node['args'] = call['args']
                        children.append(child_node)
                        
            return {'name': current_name, 'children': children}
        
        # 如果节点不在重要节点集合中且不是根节点，则将其过滤掉
        if not filter_modes or filter_mode == "none" or original_node in important_nodes:
            # 递归构建子节点
            children = []
            if original_node in relationships:
                for call in relationships[original_node]:
                    # Skip excluded functions
                    if exclude_patterns and resolve_alias_chain(call['func'], aliases) in exclude_patterns:
                        continue
                    child_node = _build(call['func'])
                    if child_node is not None:
                        child_node['args'] = call['args']
                        children.append(child_node)
            
            return {'name': current_name, 'children': children}
        else:
            # This node should be filtered out
            return None

    # 确保根节点始终存在（即使没有子节点）
    tree = _build(root)
    return tree if tree is not None else {'name': root, 'children': []}





def build_caller_tree(target: str, relationships: Dict[str, List[Dict]], 
                    relationship_labels: Dict[str, List[Dict]] = None,
                    filter_depth: Union[int, Dict[str, int]] = 2,
                    root_name: Optional[str] = None,
                    max_search_depth: int = 3) -> Dict[str, Any]:
    """
    构建调用指定函数的函数树
    
    参数:
        target: 目标函数名
        relationships: 函数调用关系
        relationship_labels: 函数标签
        filter_depth: 过滤深度，可以是整数或字典，默认为1
        root_name: 要搜索的根函数名，如果指定，将只保留通向此函数的链
        max_search_depth: 如果指定了root_name，最大搜索深度，默认为3
    """
    # 处理filter_depth参数，获取up模式的深度
    if isinstance(filter_depth, dict):
        up_depth = filter_depth.get("up", 1)  # 默认使用1
    else:
        up_depth = filter_depth
    
    # 首先构建反向调用图（被调用者 -> 调用者）
    reverse_call_graph = defaultdict(list)
    for caller, calls in relationships.items():
        caller_lower = caller.lower()  # 将调用者转换为小写
        for call in calls:
            callee = call['func'].lower()  # 确保小写一致性
            reverse_call_graph[callee].append(caller_lower)  # 存储小写版本的调用者
    
    # 调整目标函数为小写以确保一致性
    target_lower = target.lower()
    # 如果指定了root_name，也转为小写
    root_name_lower = root_name.lower() if root_name else None
    
    # 跟踪已访问的节点
    visited = set()
    
    # 创建一个函数名映射，将小写名称映射到原始名称（保留原始大小写）
    name_mapping = {name.lower(): name for name in relationships.keys()}
    name_mapping.update({call['func'].lower(): call['func'] for calls in relationships.values() for call in calls})
    
    def _build(node: str, depth=0, current_search_depth=0) -> Tuple[Optional[Dict[str, Any]], bool]:
        """递归构建调用树，返回节点及是否包含root_name"""
        # 检查是否是我们要找的根函数
        contains_root = node == root_name_lower if root_name_lower else False
        
        if node in visited:
            # 使用原始大小写展示节点名
            display_name = name_mapping.get(node, node)
            current_name = f'** {display_name}' if relationship_labels and node in relationship_labels else display_name
            return {'name': f'repeat## {current_name}', 'children': [], 'args': ''}, contains_root
            
        visited.add(node)
        
        # 使用原始大小写展示节点名
        display_name = name_mapping.get(node, node)
        current_name = f'** {display_name}' if relationship_labels and node in relationship_labels else display_name
        
        # 递归构建调用此节点的函数
        children = []
        child_contains_root = False
        
        # 判断是否继续搜索：
        # 1. 没有指定root_name时，使用正常的深度过滤
        # 2. 指定了root_name时，使用max_search_depth限制
        continue_search = (not root_name_lower and (depth < up_depth or up_depth == 0)) or \
                          (root_name_lower and current_search_depth < max_search_depth)
        
        if continue_search:
            for caller in reverse_call_graph.get(node, []):
                child_node, child_has_root = _build(caller, depth + 1, current_search_depth + 1)
                if child_node is not None:
                    if not root_name_lower or child_has_root:
                        children.append(child_node)
                        child_contains_root = child_contains_root or child_has_root
        
        # 如果当前节点是根节点或者有子节点包含根节点，则标记为包含根节点
        contains_root = contains_root or child_contains_root
        
        return {'name': current_name, 'children': children, 'args': ''}, contains_root
    
    # 从目标函数开始构建树
    tree, _ = _build(target_lower)
    return tree if tree is not None else {'name': target, 'children': [], 'args': ''}


